# -*- coding: utf-8 -*-
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import pearsonr, spearmanr

# ========= 配置 =========
FREQ   = 'Q'        # 'Q' 季度；若用月度则改为 'M'
METHOD = 'diff'     # 'diff' 同比差分；或 'log' 对数同比（仅正值总量适用）
EMP_PATH = "/Users/linshangjin/25CCM/NKU-C/dataset_csv/data_processed.xlsx"
POP_PATH = "/Users/linshangjin/25CCM/NKU-C/dataset_csv/人口结构年度数据.csv"
POP_COL  = "城镇人口比重"    # 年度文件中该列名，如不同请改
OUT_DIR  = "/Users/linshangjin/25CCM/NKU-C/t2/城乡人口结构"
# =======================

os.makedirs(OUT_DIR, exist_ok=True)

# ---------- 1) 读“就业率”并聚合 ----------
xl = pd.ExcelFile(EMP_PATH)
parts = []
for sh in xl.sheet_names:
    t = xl.parse(sh)
    if {"date","employment_rate"}.issubset(t.columns):
        tt = t[["date","employment_rate"]].copy()
        tt["date"] = pd.to_datetime(tt["date"])
        tt["employment_rate"] = pd.to_numeric(tt["employment_rate"], errors="coerce")
        parts.append(tt)
if not parts:
    raise ValueError("未在 data_processed.xlsx 中找到包含 date, employment_rate 的工作表")
emp = pd.concat(parts, ignore_index=True).dropna(subset=["date"]).sort_values("date").set_index("date")

# 聚合到季度/月度（就业率是比率，均值聚合更合适）
y = emp["employment_rate"].resample(FREQ).mean().rename("emp_"+FREQ)

# ---------- 2) 计算“每期年同比增长” ----------
lag = 4 if FREQ == 'Q' else 12

def yoy(series, method='diff'):
    s = series.copy()
    if method == 'log':
        # 对数同比要求全为正；否则回退到差分
        if (s <= 0).any():
            return s.diff(lag)
        return np.log(s).diff(lag)
    else:
        return s.diff(lag)

y_yoy = yoy(y, METHOD).rename(f"emp_{FREQ}_YoY")

# ---------- 3) 按年取均值得到“该年平均同比增长率” ----------
y_yoy_annual_mean = y_yoy.groupby(y_yoy.index.year).mean().rename("emp_yoy_avg_by_year").dropna()

# ---------- 4) 读取“城镇人口占比”年度数据并算其“年增长率” ----------
def read_annual(path):
    # 读入并兼容编码
    if path.lower().endswith((".xlsx",".xls")):
        df = pd.read_excel(path)
    else:
        try:
            df = pd.read_csv(path, encoding="utf-8")
        except UnicodeDecodeError:
            df = pd.read_csv(path, encoding="gbk")
    cols = set(df.columns)

    # 识别日期列：date / year / 年份
    if "date" in cols:
        df["date"] = pd.to_datetime(df["date"])
    elif "year" in cols:
        df["date"] = pd.to_datetime(df["year"].astype(int).astype(str) + "-12-31")
    elif "年份" in cols:
        df["date"] = pd.to_datetime(df["年份"].astype(int).astype(str) + "-12-31")
    else:
        raise ValueError(f"年度数据需包含 'date' 或 'year' 或 '年份' 列，实际列：{list(df.columns)}")

    # 数值化其余列
    for c in df.columns:
        if c not in ("date","year","年份"):
            df[c] = pd.to_numeric(df[c], errors="coerce")

    df = df.dropna(subset=["date"]).set_index("date").sort_index()
    A = df.resample("A-DEC").last()
    A.index = A.index.year
    return A

A = read_annual(POP_PATH)
if POP_COL not in A.columns:
    raise ValueError(f"在 {POP_PATH} 中未找到列：{POP_COL}（实际列：{list(A.columns)}）")

# 年增长率（占比/比率→同比差分；如需对数，把 METHOD_POP 改为 'log'）
METHOD_POP = 'diff'
def yoy_annual(s, method='diff'):
    s = s.copy()
    if method == 'log':
        s = s.dropna()
        if (s <= 0).any():
            return s.diff(1)
        return np.log(s).diff(1)
    else:
        return s.diff(1)

pop_yoy = yoy_annual(A[POP_COL].rename("urban_share"), METHOD_POP).rename("urban_share_yoy").dropna()

# ---------- 5) 对齐年份并做相关性 ----------
idx = y_yoy_annual_mean.index.intersection(pop_yoy.index)
emp_series = y_yoy_annual_mean.loc[idx]
pop_series = pop_yoy.loc[idx]

def corr_stats(a, b):
    if len(a) < 6:
        return np.nan, np.nan, np.nan, np.nan, len(a)
    rP, pP = pearsonr(a, b)
    rS, pS = spearmanr(a, b)
    return float(rP), float(pP), float(rS), float(pS), int(len(a))

rP, pP, rS, pS, nobs = corr_stats(emp_series, pop_series)

# ---------- 6) 保存结果与散点图 ----------
import csv
out_csv = os.path.join(OUT_DIR, "annual_avg_growth_corr_emp_vs_urban_share.csv")
with open(out_csv, "w", newline="", encoding="utf-8") as f:
    w = csv.writer(f)
    w.writerow(["FREQ", "EMP_METHOD", "POP_METHOD", "years_used", "pearson_r", "pearson_p", "spearman_r", "spearman_p"])
    w.writerow([FREQ, METHOD, METHOD_POP, nobs, rP, pP, rS, pS])

# 散点 + 拟合线
fig = plt.figure()
plt.scatter(pop_series.values, emp_series.values, s=45)
if len(idx) >= 3:
    a, b = np.polyfit(pop_series.values, emp_series.values, 1)
    xx = np.linspace(np.min(pop_series.values), np.max(pop_series.values), 100)
    yy = a*xx + b
    plt.plot(xx, yy)
plt.xlabel("Urban population share YoY (annual)")
plt.ylabel(f"Employment YoY avg by year ({FREQ}-based)")
plt.title("Correlation: Annual avg Employment YoY vs Urban-share YoY")
plt.tight_layout()
out_png = os.path.join(OUT_DIR, "scatter_emp_avgYoY_vs_urbanShareYoY.png")
plt.savefig(out_png, dpi=150)
plt.close(fig)

print("✅ Done")
print("对齐年份数 n =", nobs)
print("Pearson r=%.4f (p=%.4g); Spearman r=%.4f (p=%.4g)" % (rP, pP, rS, pS))
print("结果表：", out_csv)
print("散点图：", out_png)
